package com.lxg.ai.controller;

import cn.hutool.core.lang.UUID;
import com.lxg.ai.service.RagChatAssistant;
import com.lxg.ai.splitter.MyDocumentByLineSplitter;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.parser.apache.poi.ApachePoiDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.document.splitter.DocumentByRegexSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.IngestionResult;
import dev.langchain4j.store.embedding.qdrant.QdrantEmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.IOException;
import java.util.List;


@Slf4j
@RestController
@RequestMapping("/rag")
public class RagController {

//    public static void main(String[] args) {
//        String s= """
//                沃日哦\n\n
//                我儿豁""";
//        String[] split = s.split("\\s*(?>\\R)\\s*(?>\\R)\\s*");
//        System.out.println(split);
//    }


    private static List<String> POI_DOC_TYPES = List.of("doc", "docx", "ppt", "pptx", "xls", "xlsx");

    @Autowired
    private EmbeddingModel embeddingModel;

    @Autowired
    private RagChatAssistant ragChatAssistant;

    @GetMapping("ragChat")
    public void ragChat(@RequestParam String question) {
        String chat = ragChatAssistant.chat(question);
        System.out.println(chat);
    }


    @GetMapping("ragStoreChat")
    public void ragStoreChat(@RequestParam Long userId,@RequestParam String question) {
        String chat = ragChatAssistant.chat(userId,question);
        System.out.println(chat);
    }

    @PostMapping("uploadKnowledgeBase")
    public void uploadKnowledgeBase(MultipartFile file) {
        String originalFilename = file.getOriginalFilename();
        String uuid = UUID.fastUUID().toString();
        log.info("上传文档的UUID为：{}", uuid);
        String fileName = StringUtils.cleanPath(originalFilename);
        //获取文件扩展名
        String ext = getFileExtension(fileName);
        String pathName = ClassLoader.getSystemClassLoader().getResource("static").getPath() + uuid + "." + ext;
        try {
            file.transferTo(new File(pathName));
        } catch (IOException e) {
            log.error("save to local error", e);
        }
        //转化文档格式
        Document doc = transferExt(ext, pathName);
        QdrantEmbeddingStore embeddingStore = QdrantEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6334)
                .collectionName("测试向量集合")
                .build();
        //使用InMemoryEmbeddingStore嵌入
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                // 使用指定的 DocumentTransformer 转换文档。
                .documentTransformer(document -> {
                            //放置元数据uuid，进行标识
                            document.metadata().put("uuid", uuid);
                            return document;
                        }
//                new HtmlToTextDocumentTransformer()HtmlToTextDocumentTransformer in the langchain4j-document-transformer-jsoup module
                )
//        文档拆分的方案langchain4j中提供了6种
//        DocumentByCharacterSplitter 基于字符，逐个字符（含空白字符）分割
//        DocumentByRegexSplitter 基于正则的，按照自定义正则表达式分隔
//        DocumentByParagraphSplitter 基于段落的，按照连续的两个换行符（\n\n）分割
//        DocumentBySentenceSplitter 基于句子的（使用Apache OpenNLP，只支持英文，所以可以忽略）
//        DocumentByLineSplitter 基于行的，按照换行符（\n）分割
//        DocumentByWordSplitter 基于字的，将文本按照空白字符分割
                // 使用指定的 DocumentSplitter 将文档拆分为 TextSegments，1000是segment的最大token,200是上下文重叠区间
//                .documentSplitter(DocumentSplitters.recursive(1000, 200, new OpenAiTokenizer()))
//                .documentSplitter(new DocumentByParagraphSplitter(1000, 200))  试了几个分割 都不太适合 我的想法 自己写了一个
                .documentSplitter(new MyDocumentByLineSplitter(1000, 10)) //根据规则即可 不进行融合
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                .build();
        IngestionResult ingest = ingestor.ingest(doc);
        System.out.println(ingest);
    }


    //转化文档格式方法
    private static Document transferExt(String ext, String pathName) {
        Document document;
        if (ext.equalsIgnoreCase("txt")) {
            document = FileSystemDocumentLoader.loadDocument(pathName, new TextDocumentParser());
        } else if (ext.equalsIgnoreCase("pdf")) {
            document = FileSystemDocumentLoader.loadDocument(pathName, new ApachePdfBoxDocumentParser());
        } else if (POI_DOC_TYPES.contains(ext)) {
            /**
             * 使用Apache POI库将Microsoft Office文件解析为文档。这个解析器支持各种文件格式
             * ，包括doc、docx、ppt、pptx、xls和xlsx。有关支持格式的详细信息，请参阅Apache POI官方网站。
             */
            document = FileSystemDocumentLoader.loadDocument(pathName, new ApachePoiDocumentParser());
        } else {
            log.warn("该文件类型:{}无法解析，忽略", ext);
            throw new RuntimeException("该文件类型无法解析");
        }
        return document;
    }


    private static String getFileExtension(String fileName) {
        if (fileName.lastIndexOf(".") != -1 && fileName.lastIndexOf(".") != 0)
            return fileName.substring(fileName.lastIndexOf(".") + 1);
        else return "";
    }

}
